package com.aiwiown.snackmq.api.impl;

import com.aiwiown.snackmq.api.config.ConsumerConfig;
import com.aiwiown.snackmq.api.exception.SnackMQClientException;
import com.aiwiown.snackmq.common.auth.AuthenticatedUser;
import com.aiwiown.snackmq.common.dto.TopicInfo;
import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.common.message.MessageStatus;
import com.aiwiown.snackmq.common.message.MessageType;
import com.aiwiown.snackmq.common.serialization.JsonSerializer;
import com.aiwiown.snackmq.common.serialization.SerializerFactory;
import com.aiwiown.snackmq.network.client.NettyClient;
import com.aiwiown.snackmq.network.config.NetworkConfig;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import java.io.Closeable;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 【新增】消费者委托类。
 * 封装了 Push 和 Pull 消费者共享的所有核心逻辑，包括：
 * - 连接池管理
 * - 认证与授权令牌管理
 * - 主题元数据缓存与获取
 * - 向特定分区 Leader 发送请求的路由逻辑
 * <p>
 * 此类的引入遵循“组合优于继承”原则，旨在消除 AbstractSnackMQConsumer 基类。
 */
@Slf4j
public final class ConsumerDelegate implements Closeable {

    @Getter
    private final ConsumerConfig consumerConfig;
    @Getter
    private final String bootstrapBrokerAddress;
    private final Map<String, NettyClient> clientPool = new ConcurrentHashMap<>();
    private final Map<String, TopicInfo> topicInfoCache = new ConcurrentHashMap<>();
    private final JsonSerializer serializer = SerializerFactory.getJsonSerializer();
    private final AtomicBoolean started = new AtomicBoolean(false);
    private String authToken;

    public ConsumerDelegate(String bootstrapBrokerAddress, ConsumerConfig consumerConfig) {
        this.bootstrapBrokerAddress = bootstrapBrokerAddress;
        this.consumerConfig = consumerConfig;
    }

    /**
     * 启动委托类，建立初始连接并进行认证。
     *
     * @throws SnackMQClientException 如果启动失败。
     */
    public void start() throws SnackMQClientException {
        if (started.compareAndSet(false, true)) {
            log.info("Starting ConsumerDelegate for group '{}'...", consumerConfig.getConsumerGroup());
            getOrCreateClient(this.bootstrapBrokerAddress); // 建立初始连接
            authenticateAndStoreToken();
        }
    }

    /**
     * 关闭委托类，释放所有网络连接资源。
     */
    @Override
    public void close() {
        if (started.compareAndSet(true, false)) {
            log.info("Closing ConsumerDelegate for group '{}'...", consumerConfig.getConsumerGroup());
            clientPool.values().parallelStream().forEach(NettyClient::shutdown);
            clientPool.clear();
            log.info("ConsumerDelegate for group '{}' closed.", consumerConfig.getConsumerGroup());
        }
    }

    public boolean isStarted() {
        return started.get();
    }

    /**
     * 从连接池中获取或创建一个到指定 Broker 地址的 NettyClient。
     *
     * @param brokerAddress 目标 Broker 地址
     * @return 一个活跃的 NettyClient 实例
     * @throws SnackMQClientException 如果连接失败
     */
    public NettyClient getOrCreateClient(String brokerAddress) throws SnackMQClientException {
        if (!started.get()) {
            throw new SnackMQClientException("Delegate has not been started.");
        }
        NettyClient client = clientPool.get(brokerAddress);
        if (client != null && client.isActive()) {
            return client;
        }
        synchronized (clientPool) {
            client = clientPool.get(brokerAddress);
            if (client != null && client.isActive()) {
                return client;
            }
            log.info("No active client for broker [{}]. Creating a new connection.", brokerAddress);
            String[] parts = brokerAddress.split(":");
            NetworkConfig networkConfig = new NetworkConfig();
            networkConfig.setHost(parts[0]);
            networkConfig.setPort(Integer.parseInt(parts[1]));
            NettyClient newClient = new NettyClient(networkConfig);
            try {
                newClient.connect();
                clientPool.put(brokerAddress, newClient);
                return newClient;
            } catch (Exception e) {
                newClient.shutdown();
                throw new SnackMQClientException("Failed to connect to broker: " + brokerAddress, e);
            }
        }
    }

    /**
     * 向指定分区的 Leader 发送请求。
     *
     * @param request 要发送的消息
     * @param tp      目标主题分区
     * @return 一个包含 Broker 响应的 CompletableFuture
     */
    public CompletableFuture<Message> sendRequestToPartitionLeader(Message request, TopicPartition tp) {
        return fetchTopicMetadata(tp.getTopic()).thenCompose(topicInfo -> {
            if (topicInfo == null) {
                return failedFuture(new SnackMQClientException("Could not get metadata for topic: " + tp.getTopic()));
            }
            TopicInfo.PartitionAssignmentInfo partitionInfo = topicInfo.getPartitionAssignments().get(tp.getPartitionId());
            if (partitionInfo == null || partitionInfo.getLeader() == null) {
                return failedFuture(new SnackMQClientException("Could not find leader for partition " + tp.getPartitionId() + " of topic " + tp.getTopic()));
            }
            String leaderAddress = partitionInfo.getLeader().getAddress();
            try {
                NettyClient client = getOrCreateClient(leaderAddress);
                return client.sendRequest(addAuthToken(request));
            } catch (SnackMQClientException e) {
                return failedFuture(e);
            }
        });
    }

    /**
     * 向元数据 Leader 发送请求。
     *
     * @param request 要发送的消息
     * @return 一个包含 Broker 响应的 CompletableFuture
     */
    public CompletableFuture<Message> sendRequestToMetadataLeader(Message request) {
        try {
            NettyClient metadataClient = getOrCreateClient(this.bootstrapBrokerAddress);
            return metadataClient.sendRequest(addAuthToken(request));
        } catch (SnackMQClientException e) {
            return failedFuture(e);
        }
    }

    /**
     * 获取并缓存主题的元数据。
     *
     * @param topicName 主题名称
     * @return 一个包含 TopicInfo 的 CompletableFuture
     */
    public CompletableFuture<TopicInfo> fetchTopicMetadata(String topicName) {
        TopicInfo cachedInfo = topicInfoCache.get(topicName);
        if (cachedInfo != null) {
            return CompletableFuture.completedFuture(cachedInfo);
        }

        Message request = Message.builder()
                .type(MessageType.FETCH_TOPIC_METADATA)
                .topic(topicName)
                .build();

        return sendRequestToMetadataLeader(request).thenApply(response -> {
            if (response.getStatus() == MessageStatus.SUCCESS && response.getBody() != null) {
                TopicInfo fetchedInfo = serializer.deserialize(response.getBody(), TopicInfo.class);
                if (fetchedInfo != null) {
                    topicInfoCache.put(topicName, fetchedInfo);
                    return fetchedInfo;
                }
            }
            log.warn("Failed to fetch metadata for topic '{}': {}", topicName, response.getErrorMessage());
            return null;
        });
    }

    /**
     * 如果配置了认证信息，则向消息中添加认证令牌。
     *
     * @param message 原始消息
     * @return 带有认证令牌的新消息或原始消息
     */
    public Message addAuthToken(Message message) {
        if (this.authToken != null && message.getProperty("authToken") == null) {
            return message.toBuilder()
                    .property("authToken", this.authToken)
                    .build();
        }
        return message;
    }

    /**
     * 清除指定主题的元数据缓存。
     *
     * @param topicName 主题名称
     */
    public void invalidateTopicCache(String topicName) {
        topicInfoCache.remove(topicName);
    }

    private void authenticateAndStoreToken() throws SnackMQClientException {
        if (consumerConfig.getUsername() == null || consumerConfig.getUsername().isEmpty()) {
            log.info("Username is not configured for consumer group '{}', skipping authentication.", consumerConfig.getConsumerGroup());
            return;
        }
        log.info("Authenticating consumer for group '{}' with username '{}'...", consumerConfig.getConsumerGroup(), consumerConfig.getUsername());
        Message authRequest = Message.builder()
                .type(MessageType.AUTHENTICATE_REQUEST)
                .property("username", consumerConfig.getUsername())
                .property("password", consumerConfig.getPassword())
                .build();

        try {
            Message response = sendRequestToMetadataLeader(authRequest).get(5, TimeUnit.SECONDS);
            if (response.getStatus() == MessageStatus.SUCCESS && response.getBody() != null) {
                AuthenticatedUser user = serializer.deserialize(response.getBody(), AuthenticatedUser.class);
                if (user != null && user.getToken() != null) {
                    this.authToken = user.getToken();
                    log.info("Consumer for group '{}' authenticated successfully.", consumerConfig.getConsumerGroup());
                } else {
                    throw new SnackMQClientException("Authentication successful, but server did not return a valid token.");
                }
            } else {
                throw new SnackMQClientException("Authentication failed: " + response.getErrorMessage());
            }
        } catch (Exception e) {
            throw new SnackMQClientException("Failed to authenticate consumer.", e);
        }
    }

    /**
     * 创建一个已完成的、包含异常的 CompletableFuture。
     *
     * @param t 异常
     * @return 一个已异常完成的 CompletableFuture
     */
    public static <T> CompletableFuture<T> failedFuture(Throwable t) {
        CompletableFuture<T> future = new CompletableFuture<>();
        future.completeExceptionally(t);
        return future;
    }
}